use std::{collections::BTreeMap, fs::read_to_string};

use regex::Regex;

struct CityData {
    count: usize,
    names: BTreeMap<String, usize>,
    joint_set: Vec<usize>,
}
impl CityData {
    fn new() -> Self {
        Self {
            count: 0,
            names: BTreeMap::new(),
            joint_set: vec![1],
        }
    }
    /// use a usize to represent a city
    fn get_id(&mut self, name: String) -> usize {
        if let Some(id) = self.names.get(&name) {
            *id
        } else {
            self.count += 1;
            self.names.insert(name, self.count);
            self.joint_set.push(self.count);
            self.count
        }
    }
    /// return the symbol city of a province
    fn fa(&mut self, city: usize) -> usize {
        if self.joint_set[city] == city {
            return city;
        }
        let res = self.fa(self.joint_set[city]);
        self.joint_set[city] = res;
        res
    }
    /// mark to city in the same province (connect them)
    fn joint(&mut self, city1: usize, city2: usize) {
        let city1 = self.fa(city1);
        let city2 = self.fa(city2);
        self.joint_set[city1] = city2;
    }
    // return how many province (connected subgraph) in the joint set
    fn calc(&self) -> usize {
        let mut ans = 0;
        for (i, &x) in self.joint_set.iter().enumerate() {
            if i == x {
                ans += 1;
            }
        }
        ans
    }
}

pub fn count_provinces() -> String {
    let total = read_to_string("../../topic1/solutiont2/district.json").unwrap();
    let section_regex = Regex::new(r"\{([^\{\}]*)\}").unwrap();
    let city_regex = Regex::new("\"(.*?)\"").unwrap();

    let mut ans = String::new();

    for capt in section_regex.captures_iter(&total) {
        // for every test case
        let mut data = CityData::new();
        let section = capt.get(1).unwrap().as_str();
        for line in section.lines() {
            // union as in a province
            let mut last_city = 0;
            for capt in city_regex.captures_iter(line) {
                let city = data.get_id(capt.get(1).unwrap().as_str().to_owned());
                if last_city != 0 {
                    data.joint(last_city, city);
                }
                last_city = city;
            }
        }
        if !ans.is_empty() {
            ans += ",";
        }
        ans += &data.calc().to_string();
    }
    ans
}
